import os.path

import matplotlib.pyplot as plt
import numpy as np
from matplotlib import gridspec
from osgeo import gdal


def draw_tif(evi_downscal_path1, evi_avdi_path1, sm_path1, lai_path1, evi_downscal_path2, evi_avdi_path2, sm_path2, lai_path2):
    date1 = os.path.basename(evi_downscal_path1).split('_')[0]
    date2 = os.path.basename(evi_downscal_path2).split('_')[0]

    evi_downscal_dataset1 = gdal.Open(evi_downscal_path1)
    evi_avdi_dataset1 = gdal.Open(evi_avdi_path1)
    sm_dataset1 = gdal.Open(sm_path1)
    lai_dataset1 = gdal.Open(lai_path1)

    evi_downscal_dataset2 = gdal.Open(evi_downscal_path2)
    evi_avdi_dataset2 = gdal.Open(evi_avdi_path2)
    sm_dataset2 = gdal.Open(sm_path2)
    lai_dataset2 = gdal.Open(lai_path2)

    evi_downscal_data1 = evi_downscal_dataset1.ReadAsArray()
    evi_avdi_data1 = evi_avdi_dataset1.ReadAsArray()
    sm_data1 = sm_dataset1.ReadAsArray()
    lai_data1 = lai_dataset1.ReadAsArray().astype(np.float32)

    evi_downscal_data2 = evi_downscal_dataset2.ReadAsArray()
    evi_avdi_data2 = evi_avdi_dataset2.ReadAsArray()
    sm_data2 = sm_dataset2.ReadAsArray()
    lai_data2 = lai_dataset2.ReadAsArray().astype(np.float32)

    # 剔除填充值
    evi_downscal_data1[evi_downscal_data1 == -9999] = np.nan
    evi_avdi_data1[evi_avdi_data1 == -9999] = np.nan
    sm_data1[sm_data1 == -9999] = np.nan
    lai_data1[(lai_data1 == -9999) | (lai_data1 > 10)] = np.nan

    evi_downscal_data2[evi_downscal_data2 == -9999] = np.nan
    evi_avdi_data2[evi_avdi_data2 == -9999] = np.nan
    sm_data2[sm_data2 == -9999] = np.nan
    lai_data2[(lai_data2 == -9999) | (lai_data2 > 10)] = np.nan

    plt.figure(figsize=(8, 6))

    cci_min = 0
    cci_max = 0.6
    plt.subplot(4, 2, 1)
    plt.imshow(sm_data1, cmap='jet', vmin=cci_min, vmax=cci_max)
    plt.title('{}'.format(date1), fontdict={'size': 12, 'family': 'Times New Roman', 'weight': 'bold'})
    plt.xticks(np.linspace(0, sm_data1.shape[1] - 1, 4, endpoint=True), ['145°E', '146°E', '147°E', '148°E'])
    plt.yticks(np.linspace(0, sm_data1.shape[0] - 1, 3, endpoint=True), ['34.5°S', '35°S', '35.5°S'])

    plt.subplot(4, 2, 2)
    plt.imshow(sm_data2, cmap='jet', vmin=cci_min, vmax=cci_max)
    # cb = plt.colorbar(shrink=1, ticks=np.around(np.linspace(np.nanmin(sm_data2), np.nanmax(sm_data2), 6, endpoint=True), 2), fraction=0.05)
    cb = plt.colorbar(shrink=1, ticks=np.around(np.linspace(cci_min, cci_max, 7, endpoint=True), 1), fraction=0.05)
    cb.set_label('SM(m³/m³)', fontdict={'size': 12, 'family': 'Times New Roman', 'weight': 'bold'})
    plt.title('{}'.format(date2), fontdict={'size': 12, 'family': 'Times New Roman', 'weight': 'bold'})
    plt.xticks(np.linspace(0, sm_data2.shape[1] - 1, 4, endpoint=True), ['145°E', '146°E', '147°E', '148°E'])
    plt.yticks(np.linspace(0, sm_data2.shape[0] - 1, 3, endpoint=True), ['34.5°S', '35°S', '35.5°S'])

    plt.subplot(4, 2, 3)
    plt.imshow(evi_downscal_data1, cmap='jet')
    plt.title('{}'.format(date1), fontdict={'size': 12, 'family': 'Times New Roman', 'weight': 'bold'})
    plt.xticks(np.linspace(0, evi_downscal_data1.shape[1] - 1, 4, endpoint=True), ['145°E', '146°E', '147°E', '148°E'])
    plt.yticks(np.linspace(0, evi_downscal_data1.shape[0] - 1, 3, endpoint=True), ['34.5°S', '35°S', '35.5°S'])

    plt.subplot(4, 2, 4)
    plt.imshow(evi_downscal_data2, cmap='jet')
    cb = plt.colorbar(shrink=1, ticks=np.around(np.linspace(np.nanmin(evi_downscal_data2), np.nanmax(evi_downscal_data2), 5, endpoint=True), 2), fraction=0.05)
    cb.set_label('SM(m³/m³)', fontdict={'size': 12, 'family': 'Times New Roman', 'weight': 'bold'})
    plt.title('{}'.format(date2), fontdict={'size': 12, 'family': 'Times New Roman', 'weight': 'bold'})
    plt.xticks(np.linspace(0, evi_downscal_data2.shape[1] - 1, 4, endpoint=True), ['145°E', '146°E', '147°E', '148°E'])
    plt.yticks(np.linspace(0, evi_downscal_data2.shape[0] - 1, 3, endpoint=True), ['34.5°S', '35°S', '35.5°S'])

    plt.subplot(4, 2, 5)
    plt.imshow(evi_avdi_data1, cmap='jet')
    plt.title('{}'.format(date1), fontdict={'size': 12, 'family': 'Times New Roman', 'weight': 'bold'})
    plt.xticks(np.linspace(0, evi_avdi_data1.shape[1] - 1, 4, endpoint=True), ['145°E', '146°E', '147°E', '148°E'])
    plt.yticks(np.linspace(0, evi_avdi_data1.shape[0] - 1, 3, endpoint=True), ['34.5°S', '35°S', '35.5°S'])

    plt.subplot(4, 2, 6)
    plt.imshow(evi_avdi_data2, cmap='jet')
    cb = plt.colorbar(shrink=1, ticks=np.around(np.linspace(np.nanmin(evi_avdi_data2), np.nanmax(evi_avdi_data2), 5, endpoint=True), 2), fraction=0.05)
    cb.set_label('EVI AVDI', fontdict={'size': 12, 'family': 'Times New Roman', 'weight': 'bold'})
    plt.title('{}'.format(date2), fontdict={'size': 12, 'family': 'Times New Roman', 'weight': 'bold'})
    plt.xticks(np.linspace(0, evi_avdi_data2.shape[1] - 1, 4, endpoint=True), ['145°E', '146°E', '147°E', '148°E'])
    plt.yticks(np.linspace(0, evi_avdi_data2.shape[0] - 1, 3, endpoint=True), ['34.5°S', '35°S', '35.5°S'])

    plt.subplot(4, 2, 7)
    plt.imshow(lai_data1, cmap='jet')
    plt.title('{}'.format(date1), fontdict={'size': 12, 'family': 'Times New Roman', 'weight': 'bold'})
    plt.xticks(np.linspace(0, lai_data1.shape[1] - 1, 4, endpoint=True), ['145°E', '146°E', '147°E', '148°E'])
    plt.yticks(np.linspace(0, lai_data1.shape[0] - 1, 3, endpoint=True), ['34.5°S', '35°S', '35.5°S'])

    plt.subplot(4, 2, 8)
    plt.imshow(lai_data2, cmap='jet')
    cb = plt.colorbar(shrink=1, ticks=np.around(np.linspace(np.nanmin(lai_data2), np.nanmax(lai_data2), 5, endpoint=True), 5), fraction=0.05)
    cb.set_label('LAI', fontdict={'size': 12, 'family': 'Times New Roman', 'weight': 'bold'})
    plt.title('{}'.format(date2), fontdict={'size': 12, 'family': 'Times New Roman', 'weight': 'bold'})
    plt.xticks(np.linspace(0, lai_data2.shape[1] - 1, 4, endpoint=True), ['145°E', '146°E', '147°E', '148°E'])
    plt.yticks(np.linspace(0, lai_data2.shape[0] - 1, 3, endpoint=True), ['34.5°S', '35°S', '35.5°S'])

    plt.tight_layout()
    # plt.subplots_adjust(hspace=0.1)
    plt.show()


if __name__ == '__main__':
    evi_downscal_path1 = r"G:\test\SMAP_NEW\结果图\20150109_downscale_evi.tif"
    evi_avdi_path1 = r"G:\test\SMAP_NEW\结果图\20150109_evi_avdi.tif"
    sm_path1 = r"G:\test\SMAP_NEW\结果图\20150109_sm.tif"
    lai_path1 = r"G:\test\SMAP_NEW\结果图\Lai_20150109.tif"

    evi_downscal_path2 = r"G:\test\SMAP_NEW\结果图\20150109_downscale_evi.tif"
    evi_avdi_path2 = r"G:\test\SMAP_NEW\结果图\20150109_evi_avdi.tif"
    sm_path2 = r"G:\test\SMAP_NEW\结果图\20150109_sm.tif"
    lai_path2 = r"G:\test\SMAP_NEW\结果图\Lai_20150109.tif"
    draw_tif(evi_downscal_path1, evi_avdi_path1, sm_path1, lai_path1, evi_downscal_path2, evi_avdi_path2, sm_path2, lai_path2)
